package preprocessing;

import java.util.BitSet;
import java.util.HashMap;
import java.util.Vector;
import pac.SearchDatastructure;

public class WindowingSystem {
	
	private final int lastWindow = 79;
	private final double[][] windowBoundaries = new double[lastWindow][2];
	
	//separate snpList and snpMasks per window
	public BitSet[][] caseGenotypes; 
	public HashMap<Integer, Integer>[] globalPositionHash;
	public int[] counter;
	
	public WindowingSystem (  ) 
	{			
		//	System.out.println("Windows : ");
		
		//First 2 windows allotted to log scale for frequencies below 0.1%
		double lower = 0;
		for(int i=0; i<3; i++) {
			double upper = Math.pow(10, i-5);
			windowBoundaries[i][0] = lower;
			windowBoundaries[i][1] = upper;
			lower = windowBoundaries[i][1];
			//	System.out.println("Window " + i + " : " +  windowBoundaries[i][0] + " to " + windowBoundaries[i][1]);// + "\t Mean : " + windowMeans[i]);
		}
		                 
		for(int i=3; i<lastWindow-1; i++) {
			double upper = lower + (0.75 - 0.001)/(lastWindow - 5);
			windowBoundaries[i][0] = lower;
			windowBoundaries[i][1] = upper;
			lower = windowBoundaries[i][1];
			//	System.out.println("Window " + i + " : " +  windowBoundaries[i][0] + " to " + windowBoundaries[i][1]);// + "\t Mean : " + windowMeans[i]);
		}
		
		//last window
		windowBoundaries[lastWindow-1][0] = lower; 
		windowBoundaries[lastWindow-1][1] = 0.8; 				
		//	System.out.println("Window " + lastWindow + " : " +  windowBoundaries[lastWindow][0] + " to " + windowBoundaries[lastWindow][1]);
		
		//window-specific genomes
		counter = new int[lastWindow];
		globalPositionHash = new HashMap[lastWindow];
		for(int i=0; i<lastWindow;i++) 
			globalPositionHash[i] = new HashMap<Integer, Integer>();
		
	} 
	
	
	public static double Dom(double maf) {
		return 2*maf*(1-maf) + Rec(maf);
	}
	public static double Rec(double maf) {
		return maf*maf;
	}
		
	public double[][] getWindowBoundaries() {
		return windowBoundaries;
	}

	public int getNoOfWindows() {
		return lastWindow;
	}
		
	public double[][][] getExpectedCases(int[] coordinates, int numberOfCases) {
	
		double[][][] expectedCases = new double[2][2][2];
		double p_min, q_min, r_min;
		
		switch(coordinates.length) {
			case 1:
				p_min = windowBoundaries[coordinates[0]][0];
				expectedCases[1][1][1] = (double) numberOfCases * p_min;
				expectedCases[1][1][0] = numberOfCases- expectedCases[1][1][1];
				break;
				
			case 2:
				p_min = windowBoundaries[coordinates[0]][0]; 
				q_min = windowBoundaries[coordinates[1]][0];
				expectedCases[1][1][1] = numberOfCases * p_min * q_min;
				expectedCases[1][1][0] = numberOfCases * p_min * (1-q_min);
				expectedCases[1][0][1] = numberOfCases * (1-p_min) * q_min;
				expectedCases[1][0][0] = numberOfCases - expectedCases[1][1][1] - expectedCases[1][1][0] - expectedCases[1][0][1];
				break;
				
			default:
				p_min = windowBoundaries[coordinates[0]][0]; 
				q_min = windowBoundaries[coordinates[1]][0];
				r_min = windowBoundaries[coordinates[2]][0];
				expectedCases[1][1][1] = numberOfCases * p_min * q_min * r_min;
				expectedCases[1][1][0] = numberOfCases * p_min * q_min * (1-r_min);
				expectedCases[1][0][1] = numberOfCases * p_min * (1-q_min) * r_min;
				expectedCases[1][0][0] = numberOfCases * p_min * (1-q_min) * (1-r_min);
				expectedCases[0][1][1] = numberOfCases * (1-p_min)* q_min * r_min ;
				expectedCases[0][1][0] = numberOfCases * (1-p_min) * q_min * (1-r_min) ;
				expectedCases[0][0][1] = numberOfCases * (1-p_min) * (1-q_min) * r_min ;
				expectedCases[0][0][0] = numberOfCases - expectedCases[1][1][1] - expectedCases[1][1][0] 
				                                                                                                   - expectedCases[1][0][1] - expectedCases[1][0][0]
				                                                                                                   - expectedCases[0][1][1] - expectedCases[0][1][0] 
				                                                                                                                                            - expectedCases[0][0][1];
				break;
				
		}
		
		return expectedCases;
	}
		
	/*
	 * Input is maf (from frq file), output is dom and rec window index.
	 * returns an int[2] telling the caller where the recessive and dominant versions of this multi/single locus genotype go respectively.
	 */
	public int[] whichMask(double maf) {
		int[] domAndRecMask = {-1, -1};
		
		double recf = Rec(maf);
		double domf = Dom(maf);
		
		for(int i=0; i<windowBoundaries.length; i++) {
			if( windowBoundaries[i][0] <= recf && windowBoundaries[i][1] > recf ) 
				domAndRecMask[1] = i;
			if( windowBoundaries[i][0] <= domf && windowBoundaries[i][1] > domf ) {
				domAndRecMask[0] = i;
				break;
			}
		}
		return domAndRecMask;
	}
	
	
	/*
	 * Input is carrier counts, output is dom and rec window index.
	 * returns an int[2] telling the caller where the recessive and dominant versions of this multi/single locus genotype go respectively.
	 */
	private int[] whichMask ( int[] domAndRecCarrierCounts, int cohortSize ) {
		int[] domAndRecMask = {-1, -1};
		
		double domFreq = (double) domAndRecCarrierCounts[0]/cohortSize;
		double recFreq = (double) domAndRecCarrierCounts[1]/cohortSize;
		
		for(int i=0; i<windowBoundaries.length; i++) { 
			if( domFreq!=0 && windowBoundaries[i][0] <= domFreq && domFreq < windowBoundaries[i][1] ) {
				domAndRecMask[0] = i;
			}
		}

		for(int i=0; i<windowBoundaries.length; i++) { 
			if( recFreq!=0 && windowBoundaries[i][0] <= recFreq && recFreq < windowBoundaries[i][1] ) {
				domAndRecMask[1] = i;
			}
		}
		
		return domAndRecMask;
	}


public void makeWindowedGenomesFromData (SearchDatastructure searcher) 
{				
		int numberOfSNPs = searcher.userParams.getNumber_of_snps();		
		caseGenotypes = new BitSet[lastWindow][searcher.userParams.getNumber_of_cases()];	
		
		Vector<Integer>[] windowAssignedVariables = new Vector[lastWindow];
		for(int i=0; i<lastWindow;i++) 
			windowAssignedVariables[i] = new Vector<Integer>();
		
		//go through snp by snp and slot them into windows
		int wastedDomFeatures = 0;
		int wastedRecFeatures = 0;
		
		//System.out.println("Cannot test the following variables (SNP, mode) due to poor counts:");
		for(int snp=0; snp<numberOfSNPs; snp++) {
			
			//Fill in the info from controls
			int[] domThenRecCarriers = {0,0};
				
			for(int indiv=0; indiv<searcher.controlGenotypes.length; indiv++) {
				if(searcher.controlGenotypes[indiv].get(2*snp))
					domThenRecCarriers[0]++;
				if(searcher.controlGenotypes[indiv].get(2*snp+1))
					domThenRecCarriers[1]++;
			}

			searcher.snpList[snp].setControlDomCarriers(domThenRecCarriers[0]);
			searcher.snpList[snp].setControlRecCarriers(domThenRecCarriers[1]);
		
			//Fill in the Info From cases
			domThenRecCarriers[0] =0;	domThenRecCarriers[1] =0;	
			
			for(int indiv=0; indiv<searcher.caseGenotypes.length; indiv++) {
				if(searcher.caseGenotypes[indiv].get(2*snp))
					domThenRecCarriers[0]++;
				if(searcher.caseGenotypes[indiv].get(2*snp+1))
					domThenRecCarriers[1]++;
			}

			searcher.snpList[snp].setCaseDomCarriers(domThenRecCarriers[0]);
			searcher.snpList[snp].setCaseRecCarriers(domThenRecCarriers[1]);
			
			//since you are always sampling cases - use cases for masks
			
			int[] domThenRecIndex = whichMask(domThenRecCarriers, searcher.caseGenotypes.length);			
			if(domThenRecIndex[0] != -1) {
				windowAssignedVariables[domThenRecIndex[0]].add(2*snp);
			}
			else {
				//System.out.println("Discarding " + searcher.snpList[snp].getRsid() + ", d");
				wastedDomFeatures++;
			}
			
			if(domThenRecIndex[1] != -1)
				windowAssignedVariables[domThenRecIndex[1]].add(2*snp+1);
			else {
				//System.out.println("Discarding " + searcher.snpList[snp].getRsid() + ", r");
				wastedRecFeatures++;
			}
			
		}
		
		System.out.println("Discarded " + wastedDomFeatures + " dom variables and " + wastedRecFeatures + " rec variables (carried by 0 cases)."); 
				
		
		//make windowed genomes
		for (int window=0; window<lastWindow; window++) 
		{
			int numberOfVariables = windowAssignedVariables[window].size();
			for ( int indiv=0; indiv < searcher.userParams.getNumber_of_cases(); indiv++ ) 
			{
				caseGenotypes[window][indiv] = new BitSet(numberOfVariables);
				caseGenotypes[window][indiv].clear();
			}
			
			for(int varNo : windowAssignedVariables[window]) 	
			{
				for ( int indiv=0; indiv < searcher.userParams.getNumber_of_cases(); indiv++ )  
				{
					if ( searcher.caseGenotypes[indiv].get(varNo) ) 					
						caseGenotypes[window][indiv].set ( counter[window], counter[window]+1 );
				}
				globalPositionHash[window].put ( counter[window], varNo );
				counter[window]++;
			}
			
		}
		
	}



}
